/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package ru.staf.NetworkReportGenerator;

import java.io.IOException;
import java.util.HashMap;
import java.util.TreeSet;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import org.snmp4j.CommunityTarget;
import org.snmp4j.smi.Address;
import org.snmp4j.smi.GenericAddress;
import org.snmp4j.smi.OctetString;
import org.snmp4j.TransportMapping;
import org.snmp4j.transport.DefaultUdpTransportMapping;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.Snmp;
import org.snmp4j.event.ResponseListener;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.smi.OID;
import org.snmp4j.PDU;
import org.snmp4j.smi.VariableBinding;
import org.snmp4j.smi.Variable;

/**
 *
 * @author aezhov
 */
public class SnmpWalker implements ResponseListener {
    public static final String MAC_OID_D = ".1.3.6.1.2.1.17.4.3.1.2"; // dot1dTpFdbPort - learned MACs on ports for VLAN1
    public static final String MAC_OID_Q = ".1.3.6.1.2.1.17.7.1.2.2.1.2"; // dot1qTpFdbPort - learned MACs on ports

    CommunityTarget target = null;
    Address targetAddress = null;
    TransportMapping tmap = null;
    Snmp snmp = null;
    NmapXMLParser nmapParser = null;
    HashMap<String, String> snmpPortmap = null;
    HashMap<String, HashSet<String>> portIpMap = new HashMap<String, HashSet<String>>();
    HashMap<String, HashSet<String>> portMacMap = new HashMap<String, HashSet<String>>();

    public SnmpWalker(Switch sw, HashMap<String, String> snmpPortmap,
            NmapXMLParser nmapParser) throws java.net.UnknownHostException, java.io.IOException {
        this.snmpPortmap = snmpPortmap;
        this.nmapParser = nmapParser;
        
        System.out.println("SnmpWalker. Preparing walk for switch " + sw.name);
        System.out.println("    snmp version == " + sw.snmpVersion);
        System.out.println("    snmp community == " + sw.snmpCommunity);

        targetAddress = GenericAddress.parse("udp:" + sw.adminip + "/161");
        tmap = new DefaultUdpTransportMapping();
        snmp = new Snmp(tmap);
        target = new CommunityTarget(targetAddress, new OctetString(sw.snmpCommunity));
        target.setTimeout(5000);
        target.setRetries(3);
        int version = -1;
        if (sw.snmpVersion.equals("v1")) {
            version = SnmpConstants.version1;
        } else if (sw.snmpVersion.equals("v2c")) {
            version = SnmpConstants.version2c;
        }
        target.setVersion(version);
        tmap.listen();


        OID oid_d = new OID(MAC_OID_D);
        OID oid_q = new OID(MAC_OID_Q);
        PDU pdu_d = new PDU();
        PDU pdu_q = new PDU();
        pdu_d.add(new VariableBinding(oid_d));
        pdu_q.add(new VariableBinding(oid_q));
        pdu_d.setType(PDU.GETNEXT);
        pdu_q.setType(PDU.GETNEXT);
        HashSet<VariableBinding> result_d = new HashSet<VariableBinding>();
        HashSet<VariableBinding> result_q = new HashSet<VariableBinding>();
        boolean finish = false;

        TreeSet<String> ipSet = nmapParser.ipSet();
        NetworkBuzzer.shakeTheNet(ipSet);

        while (!finish) {
            ResponseEvent respEvent = snmp.send(pdu_q, target, tmap);
            PDU response = respEvent.getResponse();

            if (response != null) {
                for (int k = 0; k < response.size(); k++) {
                    VariableBinding vb = response.get(k);
                    if (vb.getOid().startsWith(oid_q)) {
                        result_q.add(vb);
                        pdu_q.remove(0);
                        pdu_q.add(vb);
                    } else {
                        finish = true;
                        break;
                    }
                }
            }
        }

        finish = false;

        while (!finish) {
            ResponseEvent respEvent = snmp.send(pdu_d, target, tmap);
            PDU response = respEvent.getResponse();

            if (response != null) {
                for (int k = 0; k < response.size(); k++) {
                    VariableBinding vb = response.get(k);
                    if (vb.getOid().startsWith(oid_d)) {
                        result_d.add(vb);
                        pdu_d.remove(0);
                        pdu_d.add(vb);
                    } else {
                        finish = true;
                        break;
                    }
                }
            }
        }

        Iterator<VariableBinding> it_d = result_d.iterator();
        Iterator<VariableBinding> it_q = result_q.iterator();

        processResult(it_d);
        processResult(it_q);

         // TEST OURSELF
        Iterator<String> testIt = portMacMap.keySet().iterator();
        while (testIt.hasNext()) {
            String port = testIt.next();
            Set<String> macs = portMacMap.get(port);
            System.out.println(port  + " :" + macs.size() + " :"+ macs);
        }

        testIt = portIpMap.keySet().iterator();
        while (testIt.hasNext()) {
            String port = testIt.next();
            Set<String> macs = portIpMap.get(port);
            System.out.println(port  + " :" + macs.size() + " :"+ macs);
        }
    }

    public HashMap<String, HashSet<String>> portToMacMap() {
        return portMacMap;
    }

    public HashMap<String, HashSet<String>> portToIpMap() {
        return portIpMap;
    }

    // Protected methods
    void processResult(Iterator<VariableBinding> result) {
        while (result.hasNext()) {
            VariableBinding vb = result.next();
            OID oid = vb.getOid();
            Variable var = vb.getVariable();
            int[] oidSubIds = oid.getValue();
            int portNum = var.toInt();
            String mac = oidToMac(oidSubIds);

            // Map snmp port id to snmp port name
            String portName = snmpPortmap.get(new Integer(portNum).toString());

            // Obtain ip for mac if available
            Set ipsForMac = nmapParser.getIpsForMac(mac);

            // Filling portIp and portMac maps
            HashSet<String> macs = portMacMap.get(portName);
            if (macs == null) {
                macs = new HashSet<String>();
                portMacMap.put(portName, macs);
            }
            macs.add(mac);

            HashSet<String> ips = portIpMap.get(portName);
            if (ips == null) {
                ips = new HashSet<String>();
                portIpMap.put(portName, ips);
            }
            if (ipsForMac != null)
                ips.addAll(ipsForMac);
        }
    }

    String oidToMac(int[] oid) {
        String result = "";
        int i = oid.length;

        result = String.format("%02x:%02x:%02x:%02x:%02x:%02x", oid[i-6], oid[i-5],
                oid[i-4], oid[i-3], oid[i-2], oid[i-1]);

        return result.toUpperCase();
    }

    // ResponseListener implementation
    public void onResponse(ResponseEvent event) {
        System.out.println("Received " + event.toString());
    }
}
